import fasttext
import numpy as np

class FastTextModel:
    def __init__(self, config):
        self.config = config
        self.model = None
    
    def train(self, train_file):
        """训练模型"""
        self.model = fasttext.train_supervised(
            train_file,
            lr=self.config.LEARNING_RATE,
            epoch=self.config.EPOCHS,
            wordNgrams=self.config.WORD_NGRAMS,
            verbose=2,
            minCount=self.config.MIN_COUNT,
            loss=self.config.LOSS
        )
        
    def predict(self, texts):
        """预测标签"""
        predictions = []
        for text in texts:
            labels, probs = self.model.predict(text)
            pred = labels[0].split('__')[-1]
            predictions.append(int(pred))
        return predictions 